﻿from torchvision import datasets,transforms
from torch.utils.data import DataLoader


def load_dataset(batch_size):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))    # 标准化,(r,g,b)均值为0.5，方差为0.5，即将数据压缩到[-1,1]之间
    ])
    
    
    #从torchversion加载自带数据集
    training_data = datasets.CIFAR10(
    root="../../data",    # 数据集存放路径
    train=True,    # 是否为训练集
    download=True,    # 是否下载
    transform=transform,    # 数据转换方式
    )

    train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)    # DataLoader是一个迭代器，每次返回一个batch的数据

    test_data = datasets.CIFAR10(
        root="../../data",
        train=False,
        download=True,
        transform=transform
    )

    test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)  
    return train_dataloader, test_dataloader